#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys, os
import moveit_commander
from moveit_commander import MoveGroupCommander, PlanningSceneInterface
from geometry_msgs.msg import PoseStamped, Pose
from moveit_msgs.msg import PlanningScene, ObjectColor
from math import pi
from copy import deepcopy
import tf
from tf.transformations import *


class MoveItObstaclesDemo:
    def __init__(self):

        # 初始化Python API 依赖的moveit_commanderC++系统，需放在前面
        moveit_commander.roscpp_initialize(sys.argv)

        # 初始化ROS节点，节点名为'moveit_obstacles_demo'
        rospy.init_node('moveit_obstacles_demo', anonymous=True)

        # 初始化场景对象
        scene = PlanningSceneInterface()
        rospy.sleep(1)
        scene_pub = rospy.Publisher('planning_scene',PlanningScene,queue_size=10)


        # 初始化需要控制的规划组。
        arm = moveit_commander.MoveGroupCommander('xarm')

        # 设置位置(单位：米)和姿态（单位：弧度）的允许误差
        arm.set_goal_position_tolerance(0.02)
        arm.set_goal_orientation_tolerance(0.03)

        # 当运动规划失败后，允许重新规划
        arm.allow_replanning(True)

        # 设置目标位姿为Home，Home是setup assistant中预先定义好的初始位姿的名称。
        arm.set_named_target('Home')
        arm.go()


        # 给每个物体设置一个特定的名字
        table_id = 'table'
        box_id = 'box'
        sphere_id = 'sphere'

        print "============ Press `Enter` to add a table to the planning scene ..."
        raw_input()

        # 将桌子(长方体桌面)添加到规划场景中
        table_size = [1.0, 1.2, 0.01]
        table_pose = PoseStamped()
        table_pose.header.frame_id = 'base_link'
        table_pose.pose.position.x = 0.0
        table_pose.pose.position.y = 0.0
        table_pose.pose.position.z =  -table_size[2]/2
        table_pose.pose.orientation.w = 1.0

        scene.add_box(table_id, table_pose, table_size)
        if self.wait_for_state_update(table_id,scene):
            rospy.loginfo("The table has been successfully added.")
        else:
            rospy.loginfo("Failed to add the table.")

        print " "
        print "============ Press `Enter` to add a sphere to the planning scene ..."
        raw_input()

        # 将一个球体添加到规划场景中
        sphere_radius = 0.12
        sphere_pose = PoseStamped()
        sphere_pose.header.frame_id = 'base_link'
        sphere_pose.pose.position.x = 0.3
        sphere_pose.pose.position.y = 0.2
        sphere_pose.pose.position.z = sphere_radius
        sphere_pose.pose.orientation.w = 1.0
        scene.add_sphere(sphere_id,sphere_pose, sphere_radius)
        if self.wait_for_state_update(sphere_id,scene):
            rospy.loginfo("The sphere has been successfully added.")
        else:
            rospy.loginfo("Failed to add the sphere.")

        print " "
        print "============ Press `Enter` to add a box to the planning scene ..."
        raw_input()
        # 将一个长方体添加到规划场景中
        box_size = [0.25, 0.05, 0.4]
        box_pose = PoseStamped()
        box_pose.header.frame_id = 'base_link'
        box_pose.pose.position.x = 0.3
        box_pose.pose.position.y = -0.2
        box_pose.pose.position.z = box_size[2]/2
        box_pose.pose.orientation.w = 1.0
        scene.add_box(box_id, box_pose, box_size)
        # rospy.sleep(1)
        if self.wait_for_state_update(box_id,scene):
            rospy.loginfo("The box has been successfully added.")
        else:
            rospy.loginfo("Failed to add the box.")

        print " "
        print "============ Press `Enter` to change the color of the box and the sphere  ..."
        raw_input()
        box_color = ObjectColor()
        box_color.id = box_id
        box_color.color.r = 0.8
        box_color.color.g = 0.8
        box_color.color.b = 0.0
        box_color.color.a = 1.0
        sphere_color = ObjectColor()
        sphere_color.id = sphere_id
        sphere_color.color.r = 0.8
        sphere_color.color.g = 0.0
        sphere_color.color.b = 0.9
        sphere_color.color.a = 1.0
        p = PlanningScene()
        p.is_diff = True
        p.object_colors.append(box_color)
        p.object_colors.append(sphere_color)
        scene_pub.publish(p)
        rospy.sleep(1)

        print " "
        print "============ Press `Enter` to move the arm ..."
        raw_input()
        # 设置末端执行器的目标位姿，参考坐标系为base_link
        # 位置通过xyz设置，姿态用四元数表示
        target_pose = PoseStamped()
        target_pose.header.frame_id = "base_link"
        target_pose.header.stamp = rospy.Time.now()
        # 末端位置
        target_pose.pose.position.x = 0.25
        target_pose.pose.position.y = -0.4
        target_pose.pose.position.z = 0.25
        # 末端姿态，四元数表示。通过quaternion_from_euler函数将RPY欧拉角转化为四元数
        quaternion = tf.transformations.quaternion_from_euler(0,0,-pi/4)
        rospy.loginfo("Quaternion is : ")
        rospy.loginfo(quaternion)

        target_pose.pose.orientation.x = quaternion[0]
        target_pose.pose.orientation.y = quaternion[1]
        target_pose.pose.orientation.z = quaternion[2]
        target_pose.pose.orientation.w = quaternion[3]

        # 设置机器臂当前的状态作为运动初始状态
        arm.set_start_state_to_current_state()

        # 设置机械臂终端运动的目标位姿
        end_effector_link = arm.get_end_effector_link()
        arm.set_pose_target(target_pose, end_effector_link)
        arm.go()
        rospy.sleep(1)

        # 设置目标位姿为Down。
        arm.set_named_target('Down')
        arm.go()
        rospy.sleep(1)

        arm.set_named_target('Home')
        arm.go()
        rospy.sleep(1)
        print " "
        print "============ Press `Enter` to remove the obstacles and exit the program ..."
        raw_input()
        scene.remove_world_object(table_id)
        scene.remove_world_object(sphere_id)
        scene.remove_world_object(box_id)
        rospy.sleep(1)

        # 干净地关闭moveit_commander并退出程序
        moveit_commander.roscpp_shutdown()
        moveit_commander.os._exit(0)

    # 判断是否已成功添加物体到规划场景
    def wait_for_state_update(self,obstacle_name,scene,timeout=4):
        start = rospy.get_time()
        seconds = rospy.get_time()

        # 在4s的时间内，循环判断我们想要的规划场景状态是否更新成功，若成功，返回True
        while (seconds - start < timeout) and not rospy.is_shutdown():

          # 判断传入的物体是否已经在规划场景中，即物体是否添加成功
          is_known = obstacle_name in scene.get_known_object_names()

          # 判断是否是我们想要的规划场景更新状态
          if is_known:
            return True

          # sleep0.1s，为物体添加留出时间。
          rospy.sleep(0.1)
          seconds = rospy.get_time()

        # If we exited the while loop without returning then we timed out
        return False


if __name__ == "__main__":
    try:
        MoveItObstaclesDemo()
    except rospy.ROSInterruptException:
        pass
